import os
import cv2
import numpy as np
import rasterio
from rasterio.plot import reshape_as_image
import geopandas as gpd
from PIL import Image as PILImage
from flask import current_app
import uuid

def process_satellite_image(filepath, image_id):
    """
    处理卫星影像，包括格式转换、预处理等
    
    Args:
        filepath (str): 影像文件路径
        image_id (int): 数据库中的影像ID
    
    Returns:
        bool: 处理成功返回True
    """
    # 获取文件扩展名
    file_ext = filepath.split('.')[-1].lower()
    
    # 处理不同格式的文件
    if file_ext in ['tif', 'tiff']:
        # 读取Tiff文件
        with rasterio.open(filepath) as src:
            # 读取影像数据
            image_data = src.read()
            # 转为图像格式
            image_data = reshape_as_image(image_data)
            
            # 如果是多光谱影像，选择RGB波段
            if image_data.shape[2] > 3:
                # 假设RGB波段是前三个波段
                image_data = image_data[:, :, 0:3]
            
            # 标准化到0-255范围
            if image_data.dtype != np.uint8:
                image_data = cv2.normalize(image_data, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
            
            # 保存为JPG用于模型处理
            processed_filename = f"{uuid.uuid4().hex}_processed_{image_id}.jpg"
            processed_path = os.path.join(current_app.config['UPLOAD_FOLDER'], processed_filename)
            
            # 使用OpenCV保存
            cv2.imwrite(processed_path, cv2.cvtColor(image_data, cv2.COLOR_RGB2BGR))
            
            # 更新数据库中的图像记录，将处理后的文件名也添加到字段中
            from models.database import db, Image
            image = Image.query.get(image_id)
            if image:
                image.processed = True
                db.session.commit()
            
            return True
    
    elif file_ext in ['jpg', 'jpeg', 'png']:
        # 对于已经是图像格式的文件，进行简单处理
        img = cv2.imread(filepath)
        
        # 调整大小为模型期望的输入尺寸
        img = cv2.resize(img, (512, 512))
        
        # 保存处理后的图像
        processed_filename = f"{uuid.uuid4().hex}_processed_{image_id}.jpg"
        processed_path = os.path.join(current_app.config['UPLOAD_FOLDER'], processed_filename)
        cv2.imwrite(processed_path, img)
        
        # 更新数据库
        from models.database import db, Image
        image = Image.query.get(image_id)
        if image:
            image.processed = True
            db.session.commit()
        
        return True
    
    else:
        raise ValueError(f"不支持的文件格式: {file_ext}")

def process_shp_file(filepath, image_id):
    """
    处理Shape文件，用于边界提取或参考
    
    Args:
        filepath (str): SHP文件路径或ZIP压缩包路径
        image_id (int): 数据库中的影像ID
    
    Returns:
        bool: 处理成功返回True
    """
    file_ext = filepath.split('.')[-1].lower()
    
    if file_ext == 'shp':
        # 直接读取shp文件
        try:
            gdf = gpd.read_file(filepath)
            
            # 根据需要处理边界信息
            # 这里可以根据实际需求进行更多处理
            
            # 更新数据库
            from models.database import db, Image
            image = Image.query.get(image_id)
            if image:
                image.processed = True
                db.session.commit()
            
            return True
        except Exception as e:
            raise ValueError(f"处理SHP文件时出错: {str(e)}")
    
    elif file_ext == 'zip':
        # 解压缩ZIP文件并处理内部的SHP文件
        import zipfile
        import tempfile
        
        try:
            # 创建临时目录
            with tempfile.TemporaryDirectory() as tmpdirname:
                # 解压ZIP文件
                with zipfile.ZipFile(filepath, 'r') as zip_ref:
                    zip_ref.extractall(tmpdirname)
                
                # 查找SHP文件
                shp_files = [f for f in os.listdir(tmpdirname) if f.endswith('.shp')]
                
                if not shp_files:
                    raise ValueError("ZIP文件中没有找到SHP文件")
                
                # 处理第一个SHP文件
                shp_path = os.path.join(tmpdirname, shp_files[0])
                gdf = gpd.read_file(shp_path)
                
                # 更新数据库
                from models.database import db, Image
                image = Image.query.get(image_id)
                if image:
                    image.processed = True
                    db.session.commit()
                
                return True
        
        except Exception as e:
            raise ValueError(f"处理ZIP文件时出错: {str(e)}")
    
    else:
        raise ValueError(f"不支持的文件格式: {file_ext}")

def augment_image(image, augmentation_type='all'):
    """
    对图像进行数据增强
    
    Args:
        image (numpy.ndarray): 输入图像
        augmentation_type (str): 增强类型，包括'flip', 'rotate', 'brightness', 'all'
    
    Returns:
        list: 增强后的图像列表
    """
    augmented_images = []
    
    if augmentation_type in ['flip', 'all']:
        # 水平翻转
        flip_h = cv2.flip(image, 1)
        augmented_images.append(flip_h)
        
        # 垂直翻转
        flip_v = cv2.flip(image, 0)
        augmented_images.append(flip_v)
    
    if augmentation_type in ['rotate', 'all']:
        # 旋转90度
        rotate_90 = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
        augmented_images.append(rotate_90)
        
        # 旋转180度
        rotate_180 = cv2.rotate(image, cv2.ROTATE_180)
        augmented_images.append(rotate_180)
    
    if augmentation_type in ['brightness', 'all']:
        # 提高亮度
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)
        v = cv2.add(v, 30)
        final_hsv = cv2.merge((h, s, v))
        bright = cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)
        augmented_images.append(bright)
        
        # 降低亮度
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)
        v = cv2.subtract(v, 30)
        final_hsv = cv2.merge((h, s, v))
        dark = cv2.cvtColor(final_hsv, cv2.COLOR_HSV2BGR)
        augmented_images.append(dark)
    
    return augmented_images

def perform_radiometric_calibration(image):
    """
    执行辐射定标，将DN值转换为辐射亮度
    
    Args:
        image (numpy.ndarray): 输入图像
        
    Returns:
        numpy.ndarray: 定标后的图像
    """
    # 这里实现辐射定标的具体算法
    # 注意：实际的定标需要依据传感器的具体参数
    calibrated = image.astype(np.float32)
    
    # 应用辐射校正系数（示例）
    gain = 0.1
    offset = 0.0
    calibrated = calibrated * gain + offset
    
    # 归一化到0-255
    calibrated = cv2.normalize(calibrated, None, 0, 255, cv2.NORM_MINMAX)
    calibrated = calibrated.astype(np.uint8)
    
    return calibrated

def perform_atmospheric_correction(image):
    """
    执行大气校正
    
    Args:
        image (numpy.ndarray): 输入图像
        
    Returns:
        numpy.ndarray: 校正后的图像
    """
    # 此处应实现实际的大气校正算法
    # 以下为简化示例
    corrected = image.copy().astype(np.float32)
    
    # 假设的大气校正参数
    for i in range(3):  # 对RGB通道分别处理
        channel = corrected[:, :, i]
        # 大气散射去除（简化模型）
        channel = (channel - 10) * 1.05
        channel[channel < 0] = 0
        channel[channel > 255] = 255
        corrected[:, :, i] = channel
    
    return corrected.astype(np.uint8)

def perform_geometric_correction(image, gcps=None):
    """
    执行几何校正
    
    Args:
        image (numpy.ndarray): 输入图像
        gcps (list): 地面控制点列表，每个点包含源坐标和目标坐标
        
    Returns:
        numpy.ndarray: 校正后的图像
    """
    # 如果没有提供GCPs，则返回原图
    if not gcps:
        return image
    
    # 提取源点和目标点
    src_points = np.array([gcp['source'] for gcp in gcps], dtype=np.float32)
    dst_points = np.array([gcp['target'] for gcp in gcps], dtype=np.float32)
    
    # 计算变换矩阵
    transform_matrix = cv2.getPerspectiveTransform(src_points, dst_points)
    
    # 应用变换
    height, width = image.shape[:2]
    corrected = cv2.warpPerspective(image, transform_matrix, (width, height))
    
    return corrected 